{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "name": "VGen_Demo.ipynb",
      "authorship_tag": "ABX9TyP0Mfd4jXJSXaPxgoxoy9x0",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/shailja-thakur/benchmarking-LLM-Verilog/blob/main/VGen_Demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Install requirements"
      ],
      "metadata": {
        "id": "n0RCnIqoTzry"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kzOirVX7E1Vy",
        "outputId": "6c5bf9cc-5140-436f-8b68-c8b5a54a30a3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting git+https://github.com/huggingface/transformers.git\n",
            "  Cloning https://github.com/huggingface/transformers.git to /tmp/pip-req-build-2rxx1t4t\n",
            "  Running command git clone -q https://github.com/huggingface/transformers.git /tmp/pip-req-build-2rxx1t4t\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from transformers==4.25.0.dev0) (4.13.0)\n",
            "Collecting tokenizers!=0.11.3,<0.14,>=0.11.1\n",
            "  Downloading tokenizers-0.13.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 7.6 MB 14.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers==4.25.0.dev0) (3.8.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from transformers==4.25.0.dev0) (2.23.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.7/dist-packages (from transformers==4.25.0.dev0) (1.21.6)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.7/dist-packages (from transformers==4.25.0.dev0) (6.0)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.7/dist-packages (from transformers==4.25.0.dev0) (4.64.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers==4.25.0.dev0) (2022.6.2)\n",
            "Collecting huggingface-hub<1.0,>=0.10.0\n",
            "  Downloading huggingface_hub-0.10.1-py3-none-any.whl (163 kB)\n",
            "\u001b[K     |████████████████████████████████| 163 kB 69.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.7/dist-packages (from transformers==4.25.0.dev0) (21.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.7/dist-packages (from huggingface-hub<1.0,>=0.10.0->transformers==4.25.0.dev0) (4.1.1)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging>=20.0->transformers==4.25.0.dev0) (3.0.9)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->transformers==4.25.0.dev0) (3.10.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->transformers==4.25.0.dev0) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->transformers==4.25.0.dev0) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->transformers==4.25.0.dev0) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->transformers==4.25.0.dev0) (2022.9.24)\n",
            "Building wheels for collected packages: transformers\n",
            "  Building wheel for transformers (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for transformers: filename=transformers-4.25.0.dev0-py3-none-any.whl size=5700835 sha256=2bc1ef10e63590cf617bf7a821b41d4311bb8ac7ed34d8596a789aa62a267ccb\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-uvfbddun/wheels/90/a5/44/6bcd83827c8a60628c5ad602f429cd5076bcce5f2a90054947\n",
            "Successfully built transformers\n",
            "Installing collected packages: tokenizers, huggingface-hub, transformers\n",
            "Successfully installed huggingface-hub-0.10.1 tokenizers-0.13.2 transformers-4.25.0.dev0\n"
          ]
        }
      ],
      "source": [
        "\n",
        "!pip install git+https://github.com/huggingface/transformers.git\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LpQNE4YYPFFQ",
        "outputId": "d7c745f3-a80b-40d7-d590-2ff42eb3ec27"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tue Nov 15 21:51:48 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   34C    P8    11W /  70W |      0MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n"
      ],
      "metadata": {
        "id": "ZiFOeqSoFKnv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load model and tokenizer"
      ],
      "metadata": {
        "id": "O_SzFmPLT6rQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "if torch.cuda.is_available():\n",
        "    device = torch.device(\"cuda\")\n",
        "    print(\"Using GPU\")\n",
        "else:\n",
        "    device = torch.device(\"cpu\")\n",
        "    print(\"Using CPU\")\n",
        "\n",
        "\n",
        "model_name = \"shailja/CodeGen_2B_Verilog\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"shailja/fine-tuned-codegen-2B-Verilog\")\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(\"shailja/fine-tuned-codegen-2B-Verilog\").to(device)\n",
        "\n"
      ],
      "metadata": {
        "id": "H4n4wD_KE3hH",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5e40f6d7-406b-4112-f26b-cf3a45ba019b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using GPU\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Try out the model for sampling Verilog"
      ],
      "metadata": {
        "id": "ijO_R0ZvUCS3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "prompt = \"//module half adder \"\n",
        "# prompt=\"// Design a 2-to-1 multiplexer.\\n module mux( \\n    input [4:0] a, b,\\n    input sel,\\n    output [4:0] out );\\n // When sel=0, assign a to out. \\n// When sel=1, assign b to out.\"\n",
        "n_steps = 1\n",
        "n = 5\n",
        "end_pattern='endmodule'\n",
        "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids.to(device)\n",
        "# Sample\n",
        "sample = model.generate(input_ids, max_length=128, temperature=0.5, top_p=0.9)\n",
        "print(tokenizer.decode(sample[0], truncate_before_pattern=[r\"endmodule\"]) + end_pattern)\n"
      ],
      "metadata": {
        "id": "_rb8UktIFJn0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Try out the model for sampling Verilog per token basis"
      ],
      "metadata": {
        "id": "m79GCcRVUFgS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "iteration = dict()\n",
        "with torch.no_grad():\n",
        "    while True:\n",
        "        \n",
        "        iteration[\"Input\"] = input_ids\n",
        "        \n",
        "        output = model(input_ids)\n",
        "\n",
        "        # select logits of the first batch and the last token and apply softmax\n",
        "        next_token_logits = output.logits[0, -1, :]\n",
        "        next_token_probs = torch.softmax(next_token_logits, dim=-1)\n",
        "        sorted_ids = torch.argsort(next_token_probs, dim=-1, descending=True)\n",
        "        # print(sorted_ids)\n",
        "        \n",
        "        \n",
        "        # store tokens with highest probability\n",
        "        for choice_idx in range(n):\n",
        "            token_id = sorted_ids[choice_idx]\n",
        "            \n",
        "            token_prob = next_token_probs[token_id].cpu().numpy()\n",
        "            # print(token_id, token_prob)\n",
        "            # print(tokenizer.decode(token_id.view(1,-1)))\n",
        "            \n",
        "            token_choice = (              \n",
        "                f\"{tokenizer.decode(token_id)} ({100*token_prob:.2f}%)\"\n",
        "            )\n",
        "            \n",
        "            iteration[f\"{choice_idx+1}\"] = token_choice\n",
        "\n",
        "        #Append predictions to list\n",
        "        input_ids = torch.cat([input_ids, sorted_ids[None,0, None]], dim=-1)\n",
        "        \n",
        "        # condition checks on stop_words detected\n",
        "        if 'endmodule' in tokenizer.batch_decode(input_ids)[0]: break\n",
        "        \n",
        "print(tokenizer.batch_decode(input_ids)[0])"
      ],
      "metadata": {
        "id": "wrtzuGMRFOVe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "1IUF4PrBFRoJ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}
